namespace Hack.Decompiling.PEA.DOS
{
  using System;
  using System.IO;
  using System.Diagnostics;
  using System.Collections.Generic;
  using System.Text;
  using System.Runtime;
  using System.Runtime.InteropServices;

  /// <summary>
  /// DOS Header
  /// </summary>
  /// <remarks>
  ///  Offset (hex) Meaning 
  ///--------------------------------------------------------------------------------
  ///  00-01  0x4d, 0x5a. This is the "magic number" of an EXE file. 
  ///         The first byte of the file is 0x4d and the second is 0x5a. 
  ///  02-03  The number of bytes in the last block of the program that are actually used. 
  ///         If this value is zero, that means the entire last block is used 
  ///         (i.e. the effective value is 512). 
  ///  04-05  Number of blocks in the file that are part of the EXE file. 
  ///         If [02-03] is non-zero, only that much of the last block is used. 
  ///  06-07  Number of relocation entries stored after the header. May be zero. 
  ///  08-09  Number of paragraphs in the header. The program's data begins just after the header, 
  ///         and this field can be used to calculate the appropriate file offset. 
  ///         The header includes the relocation entries. Note that some OSs and/or programs may fail 
  ///         if the header is not a multiple of 512 bytes. 
  ///  0A-0B  Number of paragraphs of additional memory that the program will need. 
  ///         This is the equivalent of the BSS size in a Unix program. 
  ///         The program can't be loaded if there isn't at least this much memory available to it. 
  ///  0C-0D  Maximum number of paragraphs of additional memory. 
  ///         Normally, the OS reserves all the remaining conventional memory for your program, 
  ///         but you can limit it with this field. 
  ///  0E-0F  Relative value of the stack segment. 
  ///         This value is added to the segment the program was loaded at, 
  ///         and the result is used to initialize the SS register. 
  ///  10-11  Initial value of the SP register. 
  ///  12-13  Word checksum. If set properly, the 16-bit sum of all words in the file should be zero. 
  ///         Usually, this isn't filled in. 
  ///  14-15  Initial value of the IP register. 
  ///  16-17  Initial value of the CS register, relative to the segment the program was loaded at. 
  ///  18-19  Offset of the first relocation item in the file. 
  ///  1A-1B  Overlay number. Normally zero, meaning that it's the main program. 
  /// </remarks>
  public unsafe struct DOSHeader     // 64+ (40H) Bytes
  {
    #region Statics
    public static DOSHeader it()
    {
      return new DOSHeader();
    }
    public static DOSHeader define(UInt32 bodySize, UInt32 nextHeaderAddress)
    {
      return it().Define(bodySize, nextHeaderAddress);
    }
    public static DOSHeader define(UInt32 bodySize, Byte[] stubs)
    {
      return it().Define(bodySize, stubs);
    }
    public static DOSHeader define(UInt32 bodySize, UInt32 nextHeaderAddress, Byte[] stubs)
    {
      return it().Define(bodySize, nextHeaderAddress,stubs);
    }
    public static DOSHeader define(UInt32 bodySize, Byte[] stubs, RelocationEntry[] relocations)
    {
      return it().Define(bodySize, stubs, relocations);
    }
    public static DOSHeader define(UInt32 bodySize, UInt32 nextHeaderAddress, Byte[] stubs, RelocationEntry[] relocations)
    {
      return it().Define(bodySize,nextHeaderAddress, stubs, relocations);
    }
    #endregion
    #region Consts
    public const UInt32 DefaultHeaderSize         = 0x0040;
    public const UInt16 DefaultRelocationAddress  = 0x0040;
    public const UInt16 DefaultStackSegment       = 0x0000;
    public const UInt16 DefaultStackPointer       = 0x0080;
    public const UInt16 DefaultMinParagraphs      = 0x0000;
    public const UInt16 DefaultMaxParagraphs      = 0xFFFF;
    public const UInt16 DefaultCodeSegment        = 0x0000;
    public const UInt16 DefaultCodePointer        = 0x0000;
    public const UInt16 DefaultChecksum           = 0x0000;
    public const UInt16 DefaultOverlayNumber      = 0x0000;
    public const UInt16 DefaultOEMIdentifier      = 0x0000;
    public const UInt16 DefaultOEMInformation     = 0x0000;
    public const UInt32 DefaultNextHeaderAddress  = 0x0000;

    public const UInt16 MagicValue                = 0x5A4D;
    public const UInt32 PESignature               = 0x00004550;
    public const UInt32 ParagraphSize             = 0x10;
    public const UInt32 PageSize                  = 0x200;
    public const int LenReserved1                 = 0x04;
    public const int LenReserved2                 = 0x0A;
    #endregion
    #region Fields
    #region Definition Fields
    public UInt16 MagicNumber;                          // Magic number:"MZ"
    public UInt16 PageRemainder;                        // Bytes on last page of file
    public UInt16 PageCount;                            // Pages of this file
    public UInt16 RelocationCount;                      // Relocations
    public UInt16 HeaderParagraphs;                     // Size of header in paragraphs
    public UInt16 MininumParagraphs;                    // Minimum extra paragraphs needed
    public UInt16 MaximumParagraphs;                    // Maximum extra paragraphs needed
    public UInt16 InitialStackSegment;                  // Initial (relative) SS value
    public UInt16 InitialStackPointer;                  // Initial SP value
    public UInt16 Checksum;                             // Checksum
    public UInt16 InitialCodePointer;                // Initial IP value
    public UInt16 InitialCodeSegment;                   // Initial (relative) CS value
    public UInt16 RelocationsAddress;                   // File address of relocation table:   
    public UInt16 OverlayNumber;                        // Overlay number                      //14*2 = 28 (1CH)
    public fixed UInt16 Reserveds1[LenReserved1];       // Reserved words                      //
    public UInt16 OemIdentifier;                        // OEM identifier (for e_oeminfo)      //
    public UInt16 OemInformation;                       // OEM information; e_oemid specific   //28+12 = 40 (28H)
    public fixed UInt16 Reserveds2[LenReserved2];       // Reserved words                      //40+2*10 = 60 (3CH)
    public UInt32 NextHeaderAddress;                    // File address of the new exe header  //(3CH)=PE Header (40H+:Dos stub)
    #endregion
    #region Extension Fields
    internal RelocationEntry[] relocations;
    internal Byte[] stubs;
    #endregion
    #endregion
    #region Properties
    public UInt32 FileSize
    {
      get
      {
        UInt32 size = (UInt32)this.PageCount * PageSize;

        return this.PageRemainder > 0 ? 
          (size - PageSize + this.PageRemainder) : 
          size + this.PageRemainder;

      }
      set
      {
        this.PageRemainder = (UInt16)(value % PageSize);
        this.PageCount = (UInt16)(value / PageSize);
        if (this.PageRemainder > 0)
        {
          this.PageCount++;
        }
      }
    }
    public UInt32 HeaderSize
    {
      get
      {
        return this.HeaderParagraphs * ParagraphSize;
      }
      set
      {
        this.HeaderParagraphs = (UInt16)(value / ParagraphSize);
      }
    }
    public UInt32 ImageSize
    {
      get
      {
        return this.FileSize - this.HeaderSize;
      }
      set
      {
        this.FileSize = this.HeaderSize + value;
      }
    }
    public bool HasDefinedRelocations
    {
      get
      {
        return this.relocations != null && this.relocations.Length > 0;
      }
    }
    public bool HasDefinedStubs
    {
      get
      {
        return this.stubs != null && this.stubs.Length > 0;
      }
    }
    public UInt16 DefinedRelocationCount
    {
      get
      {
        return this.HasDefinedRelocations ?
          (UInt16)this.Relocations.Length :
          (UInt16)0;
      }
    }
    public RelocationEntry[] Relocations
    {
      get
      {
        return this.relocations;
      }
    }
    public Byte[] Stubs
    {
      get
      {
        return this.stubs;
      }
    }
    #endregion
    #region Methods
    #region Publics
    public bool Read(BinaryReader reader)
    {
      UInt32 PEAddress = 0;

      return this.Read(reader, ref PEAddress);
    }
    public bool Read(BinaryReader reader, ref UInt32 PEAddress)
    {
      return this.Read(reader, ref PEAddress, true);
    }
    public bool Read(BinaryReader reader, ref UInt32 PEAddress, bool withStub)
    {
      bool valid = false;
      bool isPE = false;
      if (reader != null)
      {
        this.MagicNumber = Utils.Read16(reader);

        if (valid = (this.MagicNumber == MagicValue))
        {
          this.PageRemainder = Utils.Read16(reader);
          this.PageCount = Utils.Read16(reader);
          this.RelocationCount = Utils.Read16(reader);
          this.HeaderParagraphs = Utils.Read16(reader);
          this.MininumParagraphs = Utils.Read16(reader);
          this.MaximumParagraphs = Utils.Read16(reader);
          this.InitialStackSegment = Utils.Read16(reader);
          this.InitialStackPointer = Utils.Read16(reader);
          this.Checksum = Utils.Read16(reader);
          this.InitialCodePointer = Utils.Read16(reader);
          this.InitialCodeSegment = Utils.Read16(reader);
          this.RelocationsAddress = Utils.Read16(reader);
          this.OverlayNumber = Utils.Read16(reader);
          this.ReadReserveds1(reader);
          this.OemIdentifier = Utils.Read16(reader);
          this.OemInformation = Utils.Read16(reader);
          this.ReadReserveds2(reader);
          this.NextHeaderAddress = Utils.Read32(reader);
          this.ReadRelocations(reader);
          isPE = this.GetPEAddress(reader, ref PEAddress);
          this.FinishHeader(reader);
          //After the header, read stub of PE
          if (isPE && withStub)
          {
            this.ReadStub(reader, (PEAddress - this.HeaderSize));
          }
        }
      }
      return valid;
    }
    public void Write(BinaryWriter writer)
    {
      this.Write(writer, true);
    }
    public void Write(BinaryWriter writer,bool withStub)
    {
      if (writer != null)
      {
        Utils.Write16(writer, this.MagicNumber = MagicValue);
        Utils.Write16(writer, this.PageRemainder);
        Utils.Write16(writer, this.PageCount);
        Utils.Write16(writer, this.RelocationCount = this.DefinedRelocationCount);
        Utils.Write16(writer, this.HeaderParagraphs);
        Utils.Write16(writer, this.MininumParagraphs);
        Utils.Write16(writer, this.MaximumParagraphs);
        Utils.Write16(writer, this.InitialStackSegment);
        Utils.Write16(writer, this.InitialStackPointer);
        Utils.Write16(writer, this.Checksum);
        Utils.Write16(writer, this.InitialCodePointer);
        Utils.Write16(writer, this.InitialCodeSegment);
        Utils.Write16(writer, this.RelocationsAddress);
        Utils.Write16(writer, this.OverlayNumber);
        this.WriteReserveds1(writer);
        Utils.Write16(writer, this.OemIdentifier);
        Utils.Write16(writer, this.OemInformation);
        this.WriteReserveds2(writer);
        Utils.Write32(writer, this.NextHeaderAddress);
        this.WriteRelocations(writer);
        this.FinishHeader(writer);
        if (withStub)
        {
          this.WriteStub(writer);
        }
      }
    }
    public DOSHeader Define(UInt32 bodySize, UInt32 nextHeaderAddress)
    {
      return this.Define(bodySize,nextHeaderAddress, this.Stubs);
    }
    public DOSHeader Define(UInt32 bodySize, Byte[] stubs)
    {
      return Define(bodySize, 0, stubs);
    }
    public DOSHeader Define(UInt32 bodySize, UInt32 nextHeaderAddress, Byte[] stubs)
    {
      return this.Define(bodySize, nextHeaderAddress,stubs, this.relocations);
    }
    public DOSHeader Define(UInt32 bodySize,  Byte[] stubs, RelocationEntry[] relocations)
    {
      return this.Define(bodySize, 0, stubs, relocations);
    }
    public DOSHeader Define(UInt32 bodySize, UInt32 nextHeaderAddress, Byte[] stubs, RelocationEntry[] relocations)
    {
      this.MagicNumber              = MagicValue;

      this.HeaderSize               = DefaultHeaderSize;
      this.RelocationsAddress       = DefaultRelocationAddress;

      this.MininumParagraphs        = DefaultMinParagraphs;
      this.MaximumParagraphs        = DefaultMaxParagraphs;

      this.InitialCodeSegment       = DefaultCodeSegment;
      this.InitialCodePointer       = DefaultCodePointer;

      this.InitialStackSegment      = DefaultStackSegment;
      this.InitialStackPointer      = DefaultStackPointer;

      this.Checksum                 = DefaultChecksum;
      this.OverlayNumber            = DefaultOverlayNumber;
      this.OemIdentifier            = DefaultOEMIdentifier;
      this.OemInformation           = DefaultOEMInformation;
      this.NextHeaderAddress        = nextHeaderAddress;

      if (this.NextHeaderAddress < this.HeaderSize)
      {
        this.NextHeaderAddress = this.HeaderSize;
      }

      this.SetRelocations(relocations);

      this.SetStubs(stubs);

      if (this.HasDefinedRelocations)
      {
        this.RelocationCount = (UInt16)this.Relocations.Length;
        //Adjust header for Relocations.
        this.HeaderSize += (UInt32)this.RelocationCount * 4;

        if (this.NextHeaderAddress < this.HeaderSize)
        {
          this.NextHeaderAddress = this.HeaderSize;
        }
      }

      this.FileSize = this.HeaderSize + bodySize;

      if (this.HasDefinedStubs)
      {
        //Setup next header address for stubs.
        UInt32 stubLength = (UInt32)this.Stubs.Length;
        UInt32 stubEnding = this.HeaderSize + stubLength;

        if (this.NextHeaderAddress < stubEnding)
        {
          this.NextHeaderAddress = stubEnding;
        }

        this.FileSize += stubLength;
      }
      return this;
    }
    public DOSHeader Clone()
    {
      DOSHeader clone = DOSHeader.it();
      clone.MagicNumber = this.MagicNumber;
      clone.PageRemainder = this.PageRemainder;
      clone.PageCount = this.PageCount;
      clone.RelocationCount = this.RelocationCount;
      clone.HeaderParagraphs = this.HeaderParagraphs;
      clone.MininumParagraphs = this.MininumParagraphs;
      clone.MaximumParagraphs = this.MaximumParagraphs;
      clone.InitialStackSegment = this.InitialStackSegment;
      clone.InitialStackPointer = this.InitialStackPointer;
      clone.Checksum = this.Checksum;
      clone.InitialCodePointer = this.InitialCodePointer;
      clone.InitialCodeSegment = this.InitialCodeSegment;
      clone.RelocationsAddress = this.RelocationsAddress;
      clone.OverlayNumber = this.OverlayNumber;

      fixed (UInt16* source = this.Reserveds1)
      {
        Utils.CopyFixedArray(source, clone.Reserveds1, LenReserved1);
      }

      clone.OemIdentifier = this.OemIdentifier;
      clone.OemInformation = this.OemInformation;
      
      fixed (UInt16* source = this.Reserveds2)
      {
        Utils.CopyFixedArray(source, clone.Reserveds2, LenReserved2);
      }

      clone.NextHeaderAddress = this.NextHeaderAddress;

      clone.relocations = this.relocations.Clone() as RelocationEntry[];

      return clone;
    }
    public override string ToString()
    {
      StringBuilder builder = new StringBuilder();
      builder.AppendFormat("MagicNumber:          {0:X4}\r\n", this.MagicNumber);
      builder.AppendFormat("PageRemainder:        {0:X4}\r\n", this.PageRemainder);
      builder.AppendFormat("PageCount:            {0:X4}\r\n", this.PageCount);
      builder.AppendFormat("RelocationCount:      {0:X4}\r\n", this.RelocationCount);
      builder.AppendFormat("HeaderParagraphs:     {0:X4}\r\n", this.HeaderParagraphs);
      builder.AppendFormat("MininumParagraphs:    {0:X4}\r\n", this.MininumParagraphs);
      builder.AppendFormat("MaximumParagraphs:    {0:X4}\r\n", this.MaximumParagraphs);
      builder.AppendFormat("InitialStackSegment:  {0:X4}\r\n", this.InitialStackSegment);
      builder.AppendFormat("InitialStackPointer:  {0:X4}\r\n", this.InitialStackPointer);
      builder.AppendFormat("Checksum:             {0:X4}\r\n", this.Checksum);
      builder.AppendFormat("InitialProgramPointer:{0:X4}\r\n", this.InitialCodePointer);
      builder.AppendFormat("InitialCodeSegment:   {0:X4}\r\n", this.InitialCodeSegment);
      builder.AppendFormat("RelocationsAddress:   {0:X4}\r\n", this.RelocationsAddress);
      builder.AppendFormat("OverlayNumber:        {0:X4}\r\n", this.OverlayNumber);
      builder.AppendFormat("Reserveds1:           {0:X4}\r\n", this.FormatReserveds1());
      builder.AppendFormat("OemIdentifier:        {0:X4}\r\n", this.OemIdentifier);
      builder.AppendFormat("OemInformation:       {0:X4}\r\n", this.OemInformation);
      builder.AppendFormat("Reserveds2:           {0:X4}\r\n", this.FormatReserveds2());
      builder.AppendFormat("NextHeaderAddress:    {0:X8}\r\n", this.NextHeaderAddress);
      builder.AppendFormat("FileSize:             {0:X8}\r\n", this.FileSize);
      builder.AppendFormat("HeaderSize:           {0:X8}\r\n", this.HeaderSize);
      builder.AppendFormat("ImageSize:            {0:X8}\r\n", this.ImageSize);
      return builder.ToString();
    }
    #endregion
    #region Helpers
    internal Byte[] SetStubs(Byte[] stubs)
    {
      return this.stubs = this.NormalizeStubs(stubs);
    }
    internal RelocationEntry[] SetRelocations(RelocationEntry[] relocations)
    {
      this.RelocationCount = (this.relocations = relocations) == null ? (UInt16)0 : (UInt16)relocations.Length;

      return this.relocations;
    }
    internal Byte[] NormalizeStubs(Byte[] stubs)
    {
      return Utils.AlignBytes4(stubs);
    }
    internal bool GetPEAddress(BinaryReader reader, ref UInt32 PEAddress)
    {
      bool isPE = false;
      if (this.CanbePE())
      {
        if (this.NextHeaderAddress >= this.HeaderSize && this.NextHeaderAddress <= Utils.GetLength(reader))
        {
          Int64 currentAddress = Utils.GetPosition(reader);
          if (Utils.SetPosition(reader, this.NextHeaderAddress))
          {
            UInt32 signature = Utils.Read32(reader);
            if (isPE = (signature == PESignature))
            {
              PEAddress = this.NextHeaderAddress;
            }
          }
          Utils.SetPosition(reader, currentAddress);
        }
      }
      return isPE;
    }
    internal void ReadStub(BinaryReader reader, UInt32 stubLength)
    {
      if (reader != null)
      {
        this.stubs = Utils.ReadBytes(reader, stubLength);
      }
    }
    internal void WriteStub(BinaryWriter writer)
    {
      if (writer != null && this.HasDefinedStubs)
      {
        Utils.WriteBytes(writer, this.stubs);
      }
    }
    internal void FinishHeader(BinaryReader reader)
    {
      if (reader != null)
      {
        this.FinishHeader(reader.BaseStream);
      }
    }
    internal void FinishHeader(BinaryWriter writer)
    {
      if (writer != null)
      {
        this.FinishHeader(writer.BaseStream);
      }
    }
    internal void FinishHeader(Stream stream)
    {
      if (stream != null && stream.CanSeek)
      {
        stream.Seek((long)this.HeaderSize, SeekOrigin.Begin);
      }
    }
    internal void ReadRelocations(BinaryReader reader)
    {
      this.relocations = new RelocationEntry[this.RelocationCount];

      if (reader != null && this.RelocationCount>0)
      {

        Int64 currentAddress = Utils.GetPosition(reader);

        if (Utils.SetPosition(reader,this.RelocationsAddress))
        {
          for (UInt16 i = 0; i < this.RelocationCount; i++)
          {
            this.relocations[i] = RelocationEntry.it().Read(reader);
          }
          Utils.SetPosition(reader, currentAddress);
        }
      }
    }
    internal void WriteRelocations(BinaryWriter writer)
    {
      if (writer != null && this.HasDefinedRelocations)
      {
        Int64 currentAddress = Utils.GetPosition(writer);

        if (Utils.SetPosition(writer, this.RelocationsAddress))
        {
          for (UInt16 i = 0; i < (this.RelocationCount = this.DefinedRelocationCount); i++)
          {
            this.relocations[i].Write(writer);
          }
          Utils.SetPosition(writer, currentAddress);
        }
      }
    }
    internal void ReadReserveds1(BinaryReader reader)
    {
      fixed (UInt16* buffer = this.Reserveds1)
      {
        Utils.Read16(reader, buffer, LenReserved1);
      }
    }
    internal void ReadReserveds2(BinaryReader reader)
    {
      fixed (UInt16* buffer = this.Reserveds2)
      {
        Utils.Read16(reader, buffer, LenReserved2);
      }
    }
    internal void WriteReserveds1(BinaryWriter writer)
    {
      fixed (UInt16* buffer = this.Reserveds1)
      {
        Utils.Write16(writer, buffer, LenReserved1);
      }

    }
    internal void WriteReserveds2(BinaryWriter writer)
    {
      fixed (UInt16* buffer = this.Reserveds2)
      {
        Utils.Write16(writer, buffer, LenReserved2);
      }
    }
    internal string FormatReserveds1()
    {
      StringBuilder builder = new StringBuilder();
      fixed (UInt16* buffer = this.Reserveds1)
      {
        for (int i = 0; i < LenReserved1; i++)
        {
          builder.AppendFormat("{0:X4} ", buffer[i]);
        }
      }
      return builder.ToString(); ;
    }
    internal string FormatReserveds2()
    {
      StringBuilder builder = new StringBuilder();
      fixed (UInt16* buffer = this.Reserveds2)
      {
        for (int i = 0; i < LenReserved2; i++)
        {
          builder.AppendFormat("{0:X4} ", buffer[i]);
        }
      }
      return builder.ToString(); ;
    }

    internal bool CanbePE()
    {
      //RelocationAddress, RelocationAddress + 4 * RelocationCount;
      //3C,3D,3E,3F
      if (this.RelocationCount > 0)
      {
        //Relocation Located at NextHeaderAddress
        if ((RelocationsAddress + 4 > (DefaultHeaderSize - 4)) && RelocationsAddress < DefaultHeaderSize)
        {
          return false;
        }
      }
      return true;
    }

    #endregion
    #endregion
  }
}
